
# import magpie
import re

import jieba
import os

jieba.suggest_freq(('一个', '二维'), True)
jieba.suggest_freq(('最短'), True)
jieba.suggest_freq(('最长'), True)
jieba.suggest_freq(('波兰式'), True)


class StringDealer:

    def __init__(self):
        self.stopWordsSet = self.getStopWordsSet();
        self.dataFileBase = "data/"
        self.segmentFileBase = "segment/"

    # 获取停用词set
    def getStopWordsSet(self):
        f = open("hit_stopwords.txt",'r',encoding="utf8")
        _set = set()

        line = f.readline()
        while(line):
            _set.add(line.strip())
            line = f.readline()

        # print(_set)
        f.close()
        return _set

    # 输入一段文本, 进行分词
    def split(self, s: str)->str:
        desc = jieba.cut(s)
        _tmp = []
        for word in desc:
            if word not in self.stopWordsSet:
                _tmp.append(word)
        return " ".join(_tmp)

    # 接收一个总的文件，将其变为符合magpie的形式
    # data/huawei.txt
    def split2file(self, file_name:str):


        complete_file_name = self.dataFileBase + file_name

        f = open(complete_file_name, 'r', encoding="utf8")

        i = 1
        s = f.readline()

        while(s):
            sentence_file_name = self.segmentFileBase + file_name[:-4] + "-" + str(i) + ".txt"
            tags_file_name = self.segmentFileBase + file_name[:-4] + "-" + str(i) + ".lab"

            ss = s.split("\t")

            # 问题描述：分词后的
            desc = self.split(ss[0])

            # 问题对应的标签
            tags = ss[1]
            tags = re.sub(" ","\n",tags)


            # print(i)

            with open(sentence_file_name,"w+",encoding="utf-8") as sentence_file:
                sentence_file.write(desc)

            with open(tags_file_name,"w+",encoding="utf-8") as tags_file:
                tags_file.write(tags)

            i+=1
            s = f.readline()

        # print(ss)
        f.close()

###############################
### 把每个文件里面的每个问题拿出来
### 处理之后分成2个文件
###############################
if __name__ == '__main__':
    s = StringDealer()

    all_data = os.listdir("data")
    # print(all_data)

    for data in all_data:
        s.split2file(data)

